import numpy as np
import subprocess
import ffmpeg
from PIL import Image
import os

def load_mp4(vid_path):
    """
    使用ffmpeg加载视频文件
    Args:
        vid_path: 视频文件路径
    Returns:
        numpy array: 视频帧数组 [frames, height, width, channels]
    """
    try:
        if not os.path.exists(vid_path):
            vid_path = vid_path + '.mp4'
            if not os.path.exists(vid_path):
                raise FileNotFoundError(f"Video file not found: {vid_path}")

        probe = ffmpeg.probe(vid_path)
        video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')
        width = int(video_info['width'])
        height = int(video_info['height'])
        
        out, _ = (
            ffmpeg
            .input(vid_path)
            .output('pipe:', format='rawvideo', pix_fmt='rgb24')
            .run(capture_stdout=True, capture_stderr=True)
        )
        
        video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
        
        if video.size == 0:
            raise ValueError(f"Empty video file: {vid_path}")
            
        return video.astype('float32')
        
    except ffmpeg.Error as e:
        print(f"FFmpeg error for {vid_path}:")
        print(e.stderr.decode())
        raise
    except Exception as e:
        print(f"Error loading video {vid_path}: {str(e)}")
        raise

def load_mp4_ffmpeg(vid_path, grey=1, resolution=None):
    """
    使用ffmpeg加载并处理视频文件
    Args:
        vid_path: 视频文件路径
        grey: 是否转换为灰度图
        resolution: 目标分辨率
    Returns:
        numpy array: 处理后的视频帧数组
    """
    try:
        if not os.path.exists(vid_path):
            vid_path = vid_path + '.mp4'
            if not os.path.exists(vid_path):
                raise FileNotFoundError(f"Video file not found: {vid_path}")

        probe = ffmpeg.probe(vid_path)
        video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')
        width = int(video_info['width'])
        height = int(video_info['height'])
        
        out, _ = (
            ffmpeg
            .input(vid_path)
            .output('pipe:', format='rawvideo', pix_fmt='rgb24')
            .run(capture_stdout=True, capture_stderr=True)
        )
        
        frames = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
        
        if resolution is not None or grey:
            processed_frames = []
            for frame in frames:
                img = Image.fromarray(frame)
                if resolution:
                    img = img.resize(resolution, Image.BICUBIC)
                if grey:
                    img = img.convert('L')
                processed_frames.append(np.array(img))
            frames = np.array(processed_frames)
            
        if grey:
            frames = np.expand_dims(frames, axis=3)
            
        return frames.astype('float32')
        
    except ffmpeg.Error as e:
        print(f"FFmpeg error for {vid_path}:")
        print(e.stderr.decode())
        raise
    except Exception as e:
        print(f"Error loading video {vid_path}: {str(e)}")
        raise